[DistMuon] Drop the foreach-copy compatibility fallback - #4271
Draft
weifengpy wants to merge 2 commits into
Draft
Conversation
weifengpy
added a commit
that referenced
this pull request
Aug 21, 2026
``_foreach_copy_or_fallback_`` checks four conditions before using ``torch._foreach_copy_`` and copies pair-by-pair when any fails. Three of the four are unnecessary: the op casts between dtypes, crosses devices, and handles non-contiguous views exactly as ``Tensor.copy_`` does. Only unequal shapes are a real difference -- ``copy_`` broadcasts, the foreach op raises -- and callers pair each destination with a source built from the same ``_TensorRegion``, so their shapes always agree. When they do not, the op's own error is clearer than silently taking a slower path. Measured on Moonlight-16B-A3B, DP-shard 8 with EP 4 over 8xH100, 10 steps: the fallback ran 0 times out of 3220 multi-pair calls, and no condition failed once. Every batched call held exactly 8 uniform pairs. Rename to ``_batched_copy_``: there is no longer a fallback to name, and reusing ``_foreach_copy_`` would shadow the torch op this calls and the tests patch. Test plan: - pytest tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py -- 6 passed. The dtype test now asserts the foreach path is taken and still casts, which is the evidence for removing the guard; a new case covers the shape mismatch the op rejects. - Bit-identical numerics against the parent commit on the config above with --debug.seed 42 --debug.deterministic: step 1 loss 12.47077 grad_norm 1.1954, step 10 loss 10.69065 grad_norm 1.9209. ghstack-source-id: 3379b3e Pull-Request: #4271
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
_foreach_copy_or_fallback_checks four conditions before usingtorch._foreach_copy_and copies pair-by-pair when any fails. Three of thefour are unnecessary: the op casts between dtypes, crosses devices, and handles
non-contiguous views exactly as
Tensor.copy_does. Only unequal shapes area real difference --
copy_broadcasts, the foreach op raises -- and callerspair each destination with a source built from the same
_TensorRegion, sotheir shapes always agree. When they do not, the op's own error is clearer than
silently taking a slower path.
Measured on Moonlight-16B-A3B, DP-shard 8 with EP 4 over 8xH100, 10 steps: the
fallback ran 0 times out of 3220 multi-pair calls, and no condition failed
once. Every batched call held exactly 8 uniform pairs.
Rename to
_copy_region_views_, since there is no longer a fallback toname.
Test plan:
passed. The dtype test now asserts the foreach path is taken and still casts,
which is the evidence for removing the guard; a new case covers the shape
mismatch the op rejects.
--debug.seed 42 --debug.deterministic: step 1 loss 12.47077 grad_norm 1.1954,
step 10 loss 10.69065 grad_norm 1.9209.